///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/diis.cpp
///
/// \brief   implementation file for the diis updater
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<vector>
#include<numeric>
#include<cassert>
#include<algorithm>
#include<cmath>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/fqmc/diis.h>
#include<formic/wfn/configuration.h>
#include<formic/input/base.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/random/random.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>
#include<formic/timing/timing.h>
#include<formic/fqmc/line_search.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  constructs the DIIS updater
///
/// \param[in]      userinp        the user's input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> formic::fqmc::DIISUpdater<S>::DIISUpdater(const formic::InputBase & userinp)
{

}

template formic::fqmc::DIISUpdater<double>::DIISUpdater(const formic::InputBase & userinp);
template formic::fqmc::DIISUpdater<std::complex<double> >::DIISUpdater(const formic::InputBase & userinp);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  performs a diis update on the supplied wave function
///
/// \param[in]      userinp        the user's input options
/// \param[in]      pmtr_hist      the history of the parameters
/// \param[in]      grad_hist      the history of the gradients
/// \param[in]      sampler        object that controls sampling
/// \param[in,out]  co             a shared pointer to the value function
///                                On exit, the best update has been added to the value function
///
/// \return Whether or not the update was rejected.  Updates are rejected if they do not
///         lower the energy.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> bool formic::fqmc::DIISUpdater<S>::update(const formic::InputBase & userinp,
                                                             const std::vector<boost::shared_ptr<std::vector<double> > > & pmtr_hist,
                                                             const std::vector<boost::shared_ptr<std::vector<double> > > & grad_hist,
                                                             boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                                             const boost::shared_ptr<formic::CoeffObj<S> > & co) {

  struct local_funcs {

    static void real_to_cplx(const size_t n, const double * const r, double * const c) {
      std::copy(r, r+n, c);
    }

    static void real_to_cplx(const size_t n, const double * const r, std::complex<double> * const c) {
      for (size_t i = 0; i < n; i++)
        c[i].real() = r[i];
      for (size_t i = 0; i < n; i++)
        c[i].imag() = r[i+n];
    }

  };

  //// print a greeting
  //if (formic::mpi::rank() == 0)
  //  formic::of << "performing a DIIS update." << std::endl << std::endl;

  // ensure we are using exact sampling (diis code does not currently accomodate stochastic sampling)
  if ( !userinp.get<bool>("exact_sampling") )
    throw formic::Exception("diis update only implemented for exact sampling");

  // get the history length
  const int hist_len = userinp.get<int>("diis_hist_len");
  if ( hist_len < 1 )
    throw formic::Exception("history length must be at least one in formic::fqmc::DIISUpdater::update;  it is currently %i") % hist_len;

  // compute the B matrix
  const int bdim = hist_len + 1;
  std::vector<double> bmat(bdim*bdim, 0.0);
  for (int i = 1; i < bdim; i++) {
    bmat.at(0+bdim*i) = -1.0;
    bmat.at(i+bdim*0) = -1.0;
  }
  for (int j = 0; j < bdim-1; j++)
  for (int i = 0; i < bdim-1; i++)
    bmat.at((i+1)+bdim*(j+1)) = formic::xdotc(grad_hist.at(j)->size(), &grad_hist.at(i)->at(0), 1, &grad_hist.at(j)->at(0), 1);

  // prepare the vector for the linear equation
  std::vector<double> bvec(bdim, 0.0);
  bvec.at(0) = -1.0;

  // compute the coefficients for the linear combinations of past parameter vectors
  {
    int info = 0;
    std::vector<int> ipiv(bdim, 0);
    formic::xgesv(bdim, 1, &bmat.at(0), bdim, &ipiv.at(0), &bvec.at(0), bdim, info);
    if ( info != 0 )
      throw formic::Exception("formic::xgesv failed with info = %i in formic::fqmc::DIISUpdater::update") % info;
  }

  // compute the update
  std::vector<double> update(grad_hist.at(0)->size(), 0.0);
  for (int i = 1; i < bdim; i++) {
    if (formic::mpi::rank() == 0)
      formic::of << boost::format("bvec[%2i] = %20.12f") % i % bvec.at(i) << std::endl;
    formic::xaxpy(update.size(), bvec.at(i), &pmtr_hist.at(i-1)->at(0), 1, &update.at(0), 1);
  }
  if (formic::mpi::rank() == 0)
    formic::of << std::endl << std::endl;
  formic::xaxpy(update.size(), -1.0, &pmtr_hist.at(0)->at(0), 1, &update.at(0), 1);

  // print the dot product of the gradient and the update direction (this should be negative)
  if (formic::mpi::rank() == 0)
    formic::of << boost::format("dot of gradient with DIIS search direction = %20.6e")
                  % formic::xdotc(update.size(), &update.at(0), 1, &(**grad_hist.begin()).at(0), 1)
               << std::endl
               << std::endl;

  //formic::xaxpy(update.size(),  1.0, &pmtr_hist.at(0)->at(0), 1, &update.at(0), 1);

  // convert the update to complex form if necessary
  std::vector<S> cplx_update(co->n_der(), formic::zero(S()));
  local_funcs::real_to_cplx(co->n_der(), &update.at(0), &cplx_update.at(0));

  //co->from_vector(&cplx_update.at(0));

  // perform a line search along the update direction and remember the optimum step size
  double step_arg = userinp.get<double>("opt_diis_step_size");
  formic::fqmc::line_search(userinp, &cplx_update.at(0), userinp.get<double>("update_max_sdev"), sampler, co, step_arg);

  // return whether the update was rejected
  bool rejected = ( step_arg == 0.0 );
  return rejected;
  //return false;

}

template bool formic::fqmc::DIISUpdater<double>::update(const formic::InputBase & userinp,
                                                        const std::vector<boost::shared_ptr<std::vector<double> > > & pmtr_hist,
                                                        const std::vector<boost::shared_ptr<std::vector<double> > > & grad_hist,
                                                        boost::shared_ptr<formic::fqmc::Sampler<double> > & sampler,
                                                        const boost::shared_ptr<formic::CoeffObj<double> > & co);
template bool formic::fqmc::DIISUpdater<std::complex<double> >::update(const formic::InputBase & userinp,
                                                                       const std::vector<boost::shared_ptr<std::vector<double> > > & pmtr_hist,
                                                                       const std::vector<boost::shared_ptr<std::vector<double> > > & grad_hist,
                                                                       boost::shared_ptr<formic::fqmc::Sampler<std::complex<double> > > & sampler,
                                                                       const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & co);
